"""
股票每日历史数据抓取模块
"""
import datetime
import time
import pandas as pd
import numpy as np
from loguru import logger
import tushare as ts
import json
import os
from datetime import timedelta

from config.settings import TUSHARE_TOKEN, DATA_START_DATE, DATA_END_DATE
from database.db_engine import SessionLocal
from models.stock_models import StockBasicInfo, StockDailyData

# 设置tushare接口token
ts.set_token(TUSHARE_TOKEN)
pro = ts.pro_api()

def get_exchange_from_ts_code(ts_code):
    """
    从ts_code中获取交易所标识
    
    Args:
        ts_code: 带交易所标识的股票代码
        
    Returns:
        str: 交易所标识（SH/SZ/CY）
    """
    if ts_code.endswith('.SH'):
        return 'SH'
    elif ts_code.endswith('.SZ'):
        if ts_code.startswith('3'):
            return 'CY'
        return 'SZ'
    return None

def fetch_daily_data(trade_date=None):
    """
    获取指定日期的所有股票交易数据（包括上证、深证、创业板）并保存到数据库
    
    Args:
        trade_date: 交易日期，格式为YYYYMMDD，默认为今天
        
    Returns:
        tuple: (bool, dict) - (是否成功, 结果数据)
            结果数据格式:
            {
                'data': pandas.DataFrame,  # 股票日线数据
                'db_stats': {
                    'updated': int,  # 更新记录数
                    'inserted': int  # 新增记录数
                }
            }
    """
    try:
        # 如果未指定日期，默认获取今日数据
        if trade_date is None:
            trade_date = datetime.datetime.now().strftime("%Y%m%d")
            
        logger.info(f"开始获取 {trade_date} 的股票交易数据...")
        
        # 使用日期范围接口获取数据
        return fetch_daily_data_by_date_range(start_date=trade_date, end_date=trade_date)
            
    except Exception as e:
        logger.error(f"获取每日交易数据失败: {e}")
        return False, {'data': pd.DataFrame(), 'db_stats': {'updated': 0, 'inserted': 0}}

def fetch_daily_data_by_date_range(start_date, end_date, stock_codes=None):
    """
    获取指定日期范围和股票的交易数据（包括上证、深证、创业板）并保存到数据库
    
    Args:
        start_date: 开始日期，格式为YYYYMMDD
        end_date: 结束日期，格式为YYYYMMDD
        stock_codes: 股票代码列表，如果为空则获取所有股票数据
        
    Returns:
        tuple: (bool, dict) - (是否成功, 结果数据)
            结果数据格式:
            {
                'data': pandas.DataFrame,  # 股票日线数据
                'db_stats': {
                    'updated': int,  # 更新记录数
                    'inserted': int  # 新增记录数
                }
            }
    """
    try:
        logger.info(f"开始获取 {start_date} 至 {end_date} 的股票交易数据...")
        
        # 如果未指定股票代码，获取所有上市股票
        if not stock_codes:
            db = SessionLocal()
            stocks = db.query(StockBasicInfo.ts_code, StockBasicInfo.stock_code).filter(StockBasicInfo.status == 'L').all()
            db.close()
            stock_list = [stock[0] for stock in stocks]
            stock_code_map = {stock[0]: stock[1] for stock in stocks}  # ts_code到stock_code的映射
            logger.info(f"从数据库获取的股票列表数量: {len(stock_list)}")
        else:
            # 将股票代码转换为带交易所标识的格式
            stock_list = []
            stock_code_map = {}
            for code in stock_codes:
                if code.startswith('6'):
                    ts_code = f"{code}.SH"
                elif code.startswith('3'):
                    ts_code = f"{code}.CY"
                else:
                    ts_code = f"{code}.SZ"
                stock_list.append(ts_code)
                stock_code_map[ts_code] = code
            logger.info(f"使用指定的股票列表: {stock_list}")
        
        # 生成日期范围内的所有日期
        date_range = []
        current_date = datetime.datetime.strptime(start_date, '%Y%m%d')
        end_dt = datetime.datetime.strptime(end_date, '%Y%m%d')
        while current_date <= end_dt:
            if current_date.weekday() < 5:  # 跳过周末
                date_range.append(current_date.strftime('%Y%m%d'))
            current_date += datetime.timedelta(days=1)
        
        logger.info(f"将检查以下交易日: {date_range}")
        
        # 按交易所分组处理
        exchanges = {'SH': [], 'SZ': [], 'CY': []}
        for ts_code in stock_list:
            exchange = get_exchange_from_ts_code(ts_code)
            if exchange:
                exchanges[exchange].append(ts_code)
        
        # 分批获取数据
        batch_size = 50
        all_data = []
        
        # 按日期逐日获取数据
        for trade_date in date_range:
            logger.info(f"正在获取 {trade_date} 的数据...")
            date_data = []
            
            # 按交易所获取数据
            for exchange, exchange_stocks in exchanges.items():
                if not exchange_stocks:
                    continue
                    
                logger.info(f"正在获取{exchange}交易所的数据，共 {len(exchange_stocks)} 只股票...")
                
                for i in range(0, len(exchange_stocks), batch_size):
                    batch_stocks = exchange_stocks[i:i+batch_size]
                    stock_str = ','.join(batch_stocks)
                    
                    # 获取日线数据
                    df = pro.daily(ts_code=stock_str, trade_date=trade_date)
                    if not df.empty:
                        logger.info(f"获取到 {trade_date} {exchange} 的日线数据记录数: {len(df)}")
                    
                    # 获取每日指标数据
                    df_basic = pro.daily_basic(ts_code=stock_str, trade_date=trade_date,
                                             fields='ts_code,trade_date,turnover_rate')
                    
                    # 合并数据
                    if not df.empty and not df_basic.empty:
                        merged_df = pd.merge(df, df_basic, on=['ts_code', 'trade_date'], how='left')
                        # 添加stock_code字段
                        merged_df['stock_code'] = merged_df['ts_code'].map(stock_code_map)
                        date_data.append(merged_df)
                    elif not df.empty:
                        # 添加stock_code字段
                        df['stock_code'] = df['ts_code'].map(stock_code_map)
                        date_data.append(df)
                    
                    # 防止频繁请求
                    time.sleep(1)
            
            # 合并当天的所有批次数据
            if date_data:
                daily_result = pd.concat(date_data, ignore_index=True)
                all_data.append(daily_result)
                
                # 统计各交易所数据
                sh_count = len(daily_result[daily_result['ts_code'].str.endswith('.SH')])
                sz_count = len(daily_result[daily_result['ts_code'].str.endswith('.SZ')]) - len(daily_result[daily_result['ts_code'].str.startswith('3')])
                cy_count = len(daily_result[daily_result['ts_code'].str.startswith('3')])
                logger.info(f"{trade_date} 数据统计：上证 {sh_count} 条，深证 {sz_count} 条，创业板 {cy_count} 条")
            else:
                logger.info(f"{trade_date} 可能为非交易日，未获取到数据")
        
        # 合并所有日期的数据
        if not all_data:
            logger.warning(f"未获取到 {start_date} 至 {end_date} 的股票交易数据")
            return False, {'data': pd.DataFrame(), 'db_stats': {'updated': 0, 'inserted': 0}}
            
        result = pd.concat(all_data, ignore_index=True)
        
        # 将nan值转换为None
        result = result.replace({np.nan: None})
        
        # 记录最终数据统计
        logger.info(f"成功获取 {start_date} 至 {end_date} 的股票交易数据，共 {len(result)} 条记录")
        
        # 重命名列以匹配数据库模型
        column_mapping = {
            'pct_chg': 'pct_change',  # 涨跌幅
            'vol': 'volume',          # 成交量
            'pre_close': 'pre_close', # 昨收价
            'change': 'change',       # 涨跌额
            'ts_code': 'ts_code',     # 股票代码
            'trade_date': 'trade_date', # 交易日期
            'open': 'open',           # 开盘价
            'high': 'high',           # 最高价
            'low': 'low',             # 最低价
            'close': 'close',         # 收盘价
            'amount': 'amount',       # 成交额
            'turnover_rate': 'turnover_rate'  # 换手率
        }
        
        # 只保留数据库模型中存在的列
        result = result.rename(columns=column_mapping)
        model_columns = StockDailyData.__table__.columns.keys()
        result = result[[col for col in result.columns if col in model_columns and col != 'id']]
        
        # 转换日期格式
        result['trade_date'] = pd.to_datetime(result['trade_date'])
        
        # 保存数据到数据库
        db = SessionLocal()
        try:
            records_updated = 0
            records_inserted = 0
            
            for _, row in result.iterrows():
                # 将数据转换为字典，并处理nan值
                daily_dict = row.where(pd.notnull(row), None).to_dict()
                
                # 查询是否已存在
                exist_data = db.query(StockDailyData).filter_by(
                    ts_code=daily_dict['ts_code'],
                    trade_date=daily_dict['trade_date']
                ).first()
                
                if exist_data:
                    # 更新现有记录
                    for key, value in daily_dict.items():
                        if key in exist_data.__table__.columns.keys() and key not in ['id']:
                            setattr(exist_data, key, value)
                    records_updated += 1
                else:
                    # 创建新记录
                    new_data = StockDailyData(**daily_dict)
                    db.add(new_data)
                    records_inserted += 1
            
            # 提交事务
            db.commit()
            logger.info(f"数据保存成功：更新 {records_updated} 条记录，新增 {records_inserted} 条记录")
            
            # 转换日期为字符串格式用于API返回
            result['trade_date'] = result['trade_date'].dt.strftime('%Y%m%d')
            
            return True, {
                'data': result,
                'db_stats': {
                    'updated': records_updated,
                    'inserted': records_inserted
                }
            }
            
        except Exception as e:
            db.rollback()
            logger.error(f"保存数据到数据库失败: {e}")
            return False, {'data': pd.DataFrame(), 'db_stats': {'updated': 0, 'inserted': 0}}
        finally:
            db.close()
            
    except Exception as e:
        logger.error(f"获取日期范围内的交易数据失败: {e}")
        logger.exception("详细错误信息:")
        return False, {'data': pd.DataFrame(), 'db_stats': {'updated': 0, 'inserted': 0}}

def update_daily_data(trade_date=None):
    """
    更新指定日期的股票交易数据到数据库
    
    Args:
        trade_date: 交易日期，格式为YYYYMMDD，默认为今天
        
    Returns:
        bool: 操作是否成功
    """
    try:
        # 获取并保存日线数据
        success, result = fetch_daily_data(trade_date)
        
        if not success:
            logger.warning("获取股票日线数据失败，跳过更新")
            return False
            
        if result['data'].empty:
            logger.warning("未获取到股票日线数据，跳过更新")
            return False
            
        logger.info(f"成功更新股票日线数据，共更新 {result['db_stats']['updated']} 条记录，新增 {result['db_stats']['inserted']} 条记录")
        return True
            
    except Exception as e:
        logger.error(f"更新股票日线数据过程异常: {e}")
        return False

def update_data_from_config():
    """
    从配置文件读取日期范围并更新数据
    
    Returns:
        bool: 更新是否成功
    """
    global DATA_START_DATE
    
    # 获取日期范围
    start_date = DATA_START_DATE
    end_date = DATA_END_DATE
    
    # 验证日期
    if not start_date or not end_date:
        logger.error("未设置起始日期或结束日期")
        return False
        
    # 转换日期格式
    try:
        start_dt = datetime.datetime.strptime(start_date, "%Y%m%d")
        end_dt = datetime.datetime.strptime(end_date, "%Y%m%d")
    except ValueError as e:
        logger.error(f"日期格式错误: {e}")
        return False
        
    # 验证日期范围
    if start_dt > end_dt:
        logger.error(f"起始日期 {start_date} 大于结束日期 {end_date}")
        return False

    logger.info(f"开始更新数据，起始日期: {start_date}, 结束日期: {end_date}")
    # 更新数据
    success = update_daily_data(start_date)  # 直接使用YYYYMMDD格式
    
    #将开始日期加1天
    next_start_dt = start_dt + timedelta(days=1)
    DATA_START_DATE = next_start_dt.strftime("%Y%m%d")
    logger.info(f"开始日期更新为: {DATA_START_DATE}")
    
    return success

if __name__ == "__main__":
    # 测试代码
    update_daily_data() 